'''
Description: 
Author: suyunzheng
Date: 2021-12-07 19:45:16
LastEditTime: 2021-12-07 20:59:31
LastEditors: maple
'''
import numpy as np
import os
import sys
import torch
from core.modules.voxel_as_point import SparseGeometryFeatureEnhancement, BasicMLPBlock, BasicDeconvolutionBlock, BasicConvolutionBlock, ResidualBlock

from typing import List, Tuple
from numpy.lib import stride_tricks
from numpy.lib.twodim_base import diagflat
import torch
from torch.autograd.function import InplaceFunction
from torch.nn.modules import activation
import torchsparse.nn.functional as F
from torchsparse import PointTensor, SparseTensor
from torchsparse.nn.utils import get_kernel_offsets
from re import S
import torchsparse
import torchsparse.nn as spnn
from torch import nn
import numpy as np
from core.models.utils import initial_voxelize, point_to_voxel, voxel_to_point
# from data_utils.s3dis.s3disDataLoader import getDataLoader
from data_utils.semantickitti.SemanticKittiDataLoader import getDataLoader


class DRINet(nn.Module):
    def __init__(self, dim, **kwargs):
        super().__init__()
        init_dim = kwargs['init_dim']
        self.b1_sfe = ResidualBlock(inc=init_dim, outc=dim, ks=3, stride=1, dilation=1)
        self.b1_sgfe = SparseGeometryFeatureEnhancement(dim = dim)

        self.b2_sfe = ResidualBlock(inc=dim, outc=dim, ks=3, stride=1, dilation=1)
        self.b2_sgfe = SparseGeometryFeatureEnhancement(dim = dim)

        self.b3_sfe = ResidualBlock(inc=dim, outc=dim, ks=3, stride=1, dilation=1)
        self.b3_sgfe = SparseGeometryFeatureEnhancement(dim = dim)

        self.b4_sfe = ResidualBlock(inc=dim, outc=dim, ks=3, stride=1, dilation=1)
        self.b4_sgfe = SparseGeometryFeatureEnhancement(dim = dim)

        self.classifier = nn.Sequential(nn.Linear(dim, kwargs['num_classes']))

    def forward(self, x):
        x0 = x
        x = self.b1_sfe(x)
        x = self.b1_sgfe(x)
        x = self.b2_sfe(x)
        x = self.b2_sgfe(x)
        
        x = self.b3_sfe(x)
        x = self.b3_sgfe(x)
        x = self.b4_sfe(x)
        x = self.b4_sgfe(x)

        x = self.classifier(x.F)
        return x
        pass

if __name__ == '__main__':

    from lib.utils import Timer

    dataloader = getDataLoader(voxel_size=0.2, split='test', batch_size=2, num_points=8e4)
    iter = dataloader.__iter__()
    t = Timer()
    # model
    model = DRINet(dim=64, num_classes = 19, init_dim = 4).to('cuda:0')
    for i in range(100):
        element = iter.next()
        inputs = element['lidar'].to('cuda:0')
        targets = element['targets'].to('cuda:0')
        
        t.tic()
        outputs = model(inputs)
        print('===> iter:{}, spend time:{}'.format(i, t.toc()))
        # time = t.tic()
    pass
        
